from sklearn.linear_model import LassoCV
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score
import numpy as np


# 1. 设置超参数搜索范围
# Lasso可以自动确定最大的alpha，所以另一种设置alpha的方式是设置最小的alpha值（eps） 和 超参数的数目（n_alphas），
# 然后LassoCV对最小值和最大值之间在log域上均匀取值n_alphas个
# np.logspace(np.log10(alpha_max * eps), np.log10(alpha_max),num=n_alphas)[::-1]

def lassoTest(x_train, y_train, x_test, y_test):
    # 2 生成LassoCV实例（默认超参数搜索范围）
    lasso = LassoCV()

    # 3. 训练（内含CV）
    lasso.fit(x_train, y_train)

    # 4. 测试
    y_test_pred_lasso = lasso.predict(x_test)
    y_train_pred_lasso = lasso.predict(x_train)

    # 评估，使用r2_score评价模型在测试集和训练集上的性能
    print('The RMSE of LassoCV on test is', np.sqrt(mean_squared_error(y_test, y_test_pred_lasso)))
    print('The RMSE of LassoCV on train is', np.sqrt(mean_squared_error(y_train, y_train_pred_lasso)))
    #print('The r2 score of LassoCV on test is', r2_score(y_test, y_test_pred_lasso))
    #print('The r2 score of LassoCV on train is', r2_score(y_train, y_train_pred_lasso))
    mses = np.mean(lasso.mse_path_, axis=1)
    print('lasso.alphas_ is:', lasso.alphas_)
    return mses, lasso.alphas_, lasso.coef_
